#!/usr/bin/env python3
# record_frames.py

import argparse
import os
import time
from ctypes import addressof, c_ubyte

import gxipy as gx
from gxipy.gxidef import GxPixelFormatEntry, DxValidBit
from gxipy.ImageFormatConvert import ImageFormatConvert
import numpy as np
import cv2

def get_best_valid_bits(pixel_format):
    """映射 RawImage 像素格式到最佳有效位设置。"""
    # 这里只列示常见 MONO8 与 RGB8，对应映射可参照官方 SDK 文档
    if pixel_format in (GxPixelFormatEntry.MONO8, GxPixelFormatEntry.RGB8):
        return DxValidBit.BIT0_7
    # 其余类型请补充完整映射……
    return DxValidBit.BIT0_7

def convert_to_bgr(raw_img, img_conv):
    """
    用 ImageFormatConvert 将 RawImage 转为 RGB8，然后转为 BGR NumPy 数组。
    """
    img_conv.set_dest_format(GxPixelFormatEntry.RGB8)
    img_conv.set_valid_bits(get_best_valid_bits(raw_img.get_pixel_format()))
    buf_size = img_conv.get_buffer_size_for_conversion(raw_img)
    buf = (c_ubyte * buf_size)()
    img_conv.convert(raw_img, addressof(buf), buf_size, False)
    arr = np.frombuffer(buf, dtype=np.uint8, count=buf_size)  # 直接从缓冲区创建数组【&#8203;:contentReference[oaicite:4]{index=4}】
    h, w = raw_img.get_height(), raw_img.get_width()
    rgb = arr.reshape(h, w, 3)  # 重塑为 H×W×3【&#8203;:contentReference[oaicite:5]{index=5}】
    return cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)

def parse_args():
    """命令行参数解析：输出目录、时间戳文件、曝光、相机索引。"""
    parser = argparse.ArgumentParser(
        description="使用 GXIPY 采集帧并保存为 PNG 与时间戳")
    parser.add_argument(
        "--output-dir", required=True,
        help="保存帧图像的目录（例如 frames/）")
    parser.add_argument(
        "--output-ts", required=True,
        help="保存时间戳的 TXT 文件路径（例如 ts.txt）")
    parser.add_argument(
        "--exposure", type=float, default=70000.0,
        help="曝光时间（微秒）")
    parser.add_argument(
        "--cam-index", type=int, default=1,
        help="相机设备索引（1 表示第一个相机）")
    return parser.parse_args()

def main():
    args = parse_args()  # 获取命令行参数【&#8203;:contentReference[oaicite:6]{index=6}】

    # 创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)

    # 初始化相机
    dev_mgr = gx.DeviceManager()
    num, _ = dev_mgr.update_all_device_list()
    if num < args.cam_index:
        print(f"未找到索引为 {args.cam_index} 的相机")
        return
    cam = dev_mgr.open_device_by_index(args.cam_index)
    feat = cam.get_remote_device_feature_control()

    # 设置曝光与软件触发
    feat.get_float_feature("ExposureTime").set(args.exposure)
    feat.get_enum_feature("TriggerMode").set("On")
    feat.get_enum_feature("TriggerSource").set("Software")

    # 图像转换器
    img_conv = dev_mgr.create_image_format_convert()
    ds = cam.data_stream[0]
    trig = feat.get_command_feature("TriggerSoftware")

    # 启动采集
    cam.stream_on()
    print("<Start acquisition>")

    # 打开时间戳日志
    logf = open(args.output_ts, "w", encoding="utf-8")
    logf.write("frame_idx\ttimestamp\n")

    frame_idx = 0
    cv2.namedWindow("Camera", cv2.WINDOW_NORMAL)

    try:
        while True:
            # 触发采集
            trig.send_command()
            raw = ds.get_image()  # 读取 RawImage
            if raw:
                frame = convert_to_bgr(raw, img_conv)
                ts = time.time()
                # 保存帧为 PNG
                png_path = os.path.join(
                    args.output_dir,
                    f"frame_{frame_idx:06d}.png")
                cv2.imwrite(png_path, frame)  # 写文件【&#8203;:contentReference[oaicite:7]{index=7}】
                # 记录时间戳
                logf.write(f"{frame_idx}\t{ts:.6f}\n")
                frame_idx += 1

                # 显示
                cv2.imshow("Camera", frame)  # 预览亦可【&#8203;:contentReference[oaicite:8]{index=8}】

            # 按 q 键退出
            if cv2.waitKey(1) & 0xFF == ord("q"):
                break
            time.sleep(0.01)
    except KeyboardInterrupt:
        pass
    finally:
        # 清理资源
        logf.close()
        cam.stream_off()
        ds.unregister_capture_callback()
        cam.close_device()
        cv2.destroyAllWindows()
        print("<Stop acquisition>")

if __name__ == "__main__":
    main()
